import time

from script import snode, NodeResult, common_client, NodeTaskCtx, EvalInfo, EvalTaskCtx, PlanKeyDefine, EvalResult

#  脚本文件随意编写，那么流程怎么串呢？
# 在你需要将其设置为一个节点的方法加上装饰器（类似于java注解），系统就会识别到这是一个节点
# 节点与节点之间可以通过装饰器的accept_nodes属性进行串联，传入值是上游节点的名称，这个accept_nodes属性将单个py脚本里的所有函数串起来
# 所得链路图为多源多输出的模式，即可以有多个源头，也可以有多个尾节点
# 每个函数接受两个参数，类型为ScriptNodeData，非空，dict ，非空
# 返回值可以是NodeResult 也可以是 NodeResult数组 亦或是字典，但是单个元素的结构要和NodeResult相似
# qpm 以及 rt一起设置时系统会根据配置在跑批时自动设置并发量，支持动态更新
# accept_types 可限制接受的类型数据，假设我有个节点判断文本是否是饿么诉求，后面还会分别对有/没有的节点跑不同的任务，这时就有用了
# explain_sql 解释性sql 统一表结构单一无法反映真实节点结构，配置个sql在导出后可以查询想要的结果


from script.common_storage import ScriptNodeData


def eval1(info: EvalInfo, ctx: EvalTaskCtx) -> EvalResult:
    return EvalResult({
        "演示样本917-2": [{
            "id": 392928,
            "content": "{\"text\": \"你好坏\"}",
            "评估结果": "正确"
        }]
    })


def prepare(datas: list[ScriptNodeData], ctx: NodeTaskCtx) -> dict:
    # 每个worker启动时会
    return {"name":'1'}


@snode(desc='方法a', accept_types=['1', '2'], qpm=6, rt=2000, content_example='''{
    "text":"你要测试的样本"
}''', explain_sql='''
select content_key as 文本,
content_type as 正负向
from {table} 
where {pt}
and content_type <> 'hard'
''', eval_func='eval1', prepare_func='prepare', plan_keys=[PlanKeyDefine('template', desc='这是模板，非必填'), PlanKeyDefine('model', desc='模型，非必填')])
def funa(data: ScriptNodeData, ctx: NodeTaskCtx):
    content_json = data.content_as_json()
    if not content_json or 'text' not in content_json:
        raise ValueError('内容不合法')
    template = ctx.get_plan_config('template', '''请判断以下这行文本表达的情绪是正反馈还是负反馈
    需要判断的文本："""{text}"""
    请直接回答：正向/负向
    如果无法判断请回答：不确定''')
    model = ctx.get_plan_config('model', 'gpt-3.5-turbo')
    llm_res = common_client.llm_predict(template, content_json, model=model)
    llm_res.check_success()
    type = llm_res.match_result(['正向', '负向'], 'hard')
    return NodeResult(content_type=type, content=content_json, content_key=content_json.get('text'))


@snode(desc='方法b', accept_nodes=["funa"], rt=3000, qpm=10)
def funb(data: ScriptNodeData, ctx: NodeTaskCtx):
    time.sleep(3)
    return NodeResult(content_key=data.content_type, content=data.content)


@snode(desc='方法c', accept_nodes=["funa", "funb"])
def func(data: ScriptNodeData, ctx: NodeTaskCtx):
    return None
